import matplotlib.pyplot as plt
import tushare as ts
from mpl_finance import  candlestick2_ochl

def drawK(pos, open, close, high, low):
    if close>open: # 是上涨
        myColor='red'
        myHeight=close-open
        myBottom=open
    if close<open: # 是下跌
        myColor='green'
        myHeight=open-close
        myBottom=close
    # 根据开盘价和收盘价绘制长方形实体
    plt.bar(pos, height=myHeight, width=0.2, bottom=myBottom, color=myColor)
    # 根据最高价和最低价绘制上下影线
    plt.vlines(pos, low, high, myColor)

# 绘制Y值
drawK(0, 10.2, 10.8, 9, 11)
drawK(1, 10.8, 11.9, 10.8, 12)
drawK(2, 10.9, 11.9, 10, 12)
drawK(3, 12, 10, 12, 9)
drawK(4, 9.8, 10.2, 9, 10.5)
drawK(5, 10.8, 11.9, 10.8, 12)

# 定义时间范围
day = ['20231001','20231002','20231003','20231004','20231005','20231006']

# x轴标签
plt.xticks(range(len(day)), day)
# y轴取值范围
plt.ylim(0, 15)
# 展示
plt.show()

# tushare

# set token
ts.set_token('9084dfbcc6245ce84b632fbad1f2c24539b71aeec5f021075bfbcecf')
# init pro api
pro = ts.pro_api()
# get 000001股票的数据
kd = pro.daily(ts_code='000001.SZ', start_date='20230701', end_date='20230720')
print(kd)
# 白板
# fig=plt.figure()
# ax=fig.add_subplot
# 绘制k线
candlestick2_ochl(
    plt.axes(), 
    opens=kd['open'].values,
    closes=kd['close'].values,
    highs=kd['high'].values,
    lows=kd['low'].values,
    width=0.7,
    )
# 绘制x轴
plt.xticks(range(len(kd['trade_date'].values)), kd['trade_date'], rotation=30)
# 网格线
plt.grid(True)
# plt.ylim(10, 13)
plt.show()